import matplotlib
from scipy import shape,concatenate,nansum,cumproduct,zeros,nan,polyfit,array,isnan,nansum,copy,roll,sum,diff,sort,poly1d,nan_to_num
from pylab import find
def rollavg(data,n=5,axis=None,circular=False):
    if n<=1: return data
    datat = data.swapaxes(0,axis) if axis is not None else data
    res = copy(datat).astype('f')
    res_tmp = array(map(lambda x: roll(res,-x,axis=0),range(-n/2,-n/2+n)))
    res_tmp = nansum(res_tmp,axis=0)/sum(~isnan(res_tmp),0)
    if circular is False:
      res[(n/2):(-n/2)] = res_tmp[(n/2):(-n/2)]
    else:
      res = res_tmp
    if axis is not None:
      res=res.swapaxes(0,axis)
    return res

def reduceData(M,n,domean=True):#mymean
    #cM = M.cumsum(0)[(int(n)-1)::int(n),:]
    cM=nan_to_num(M).cumsum(0)
    cM[isnan(M)]=nan
    cM=cM[(int(n)-1)::int(n),:]
    cM[1:] = cM[1:,:] - cM[:-1,:]
    return cM/(float(n) if domean else 1.)
 
def mycorrcoef(data):
  Gmean_calc = (data.T[:]-data.mean(1)).T
  #normalize=abs(Gmean_calc).sum(1)
  #print normalize.min()
  res=zeros((len(data),len(data)))
  for i in range(len(data)):
    for j in range(len(data)-i):
      #res[i,i+j] = (Gmean_calc[i]*Gmean_calc[i+j]).sum()/(normalize[i]*normalize[i+j])
      res[i,i+j] = (Gmean_calc[i]*Gmean_calc[i+j]).mean()#/(normalize[i]*normalize[i+j])
  return res


def deldup(PE):
  if len(PE)==0:
    return []
  M = dict(zip(PE[:,0], PE[:,0])).keys()
  pe=list(PE)
  dellist=[]
  for i in M:
    dellist += list(list(PE[:,0]).index(i)+find(sum(diff(PE[PE[:,0]==i,1:],axis=0)==0,1)>0))
  for d in sort(dellist)[::-1]:
      del(pe[d])
  print "Deleted %s sets of peaks which used the same peak" % len(dellist)
  return array(pe)  


def slope_diff(V,I,nfitlength=5,dI=1):
    ltraces, ldat,lV= shape(I)
    stop = lV-nfitlength
    dIdV = [[nan]*ltraces]*lV
    i=-nfitlength/2-1
    Vt = concatenate((V[i:],V[0:i+nfitlength]))
    It = concatenate((I[:,:,i:],I[:,:,0:i+nfitlength]),axis=2)
    for i in xrange(-nfitlength/2,stop+nfitlength/2):
      Vt = concatenate((Vt[1:],[V[(i+nfitlength-1)%len(V)]]))
      It = concatenate((It[:,:,1:],I[:,:,[(i+nfitlength-1)%len(V)]]),axis=2)
      Vtt = concatenate([Vt for j in range(ldat)])
      Itt = concatenate(It.T,axis=0)
      dIdV[i+nfitlength/2] = polyfit(Vtt, Itt,dI)[0]*dI
    dIdV = array(dIdV).T
    dIdV[isnan(I).any(2)]=nan
    return dIdV

def correctbasline(v,I,pol=9,n=10,zerocor=10):
  vf = copy(v)
  If = copy(I)
  listrange = find(isnan(If)==False)
  for i in range(n):
    pf = poly1d(polyfit(array([-10]*zerocor+list(vf[listrange])),array([min(I[listrange])]*zerocor+list(If[listrange])),pol))
    If[I<pf(v)] = I[I<pf(v)]
    If[I>pf(v)] = pf(v)[I>pf(v)]
  return I-pf(v)


def intersect(a, b):
  """ return the intersection of two lists """
  return list(set(a) & set(b))

cdict = {'red': ((0.0, 1.0, 1.0),
                 (0.1, 0.0, 0.0),
                 (.8, 1.0, 1.0),
                 (1, .0, .0)),
         'green': ((0.0, .90, .90),
                   (.8, 0.0, 0.0),
                   (1, .0, .0)),
         'blue': ((0.0, .90, .90),
                  (0.2, .90, .90),
                  (.8, 0.0, 0.0),
                  (1, .0, .0))
                }
my_cmap = matplotlib.colors.LinearSegmentedColormap('my_colormap',cdict,256)

nanmean = lambda v,axis=None: nansum(v,axis=axis)/((cumproduct(shape(v))[-1] if axis is None else shape(v)[axis])-sum(isnan(v),axis=axis))
__nanstd__ = lambda v,axis=None:nanmean(v,axis=axis).reshape([(1 if i==axis else s) for i,s in enumerate(shape(v))])-v
nanstd = lambda v,axis=None: sqrt(nanmean((nanmean(v)-v if axis is None else __nanstd__(v,axis))**2,axis=axis))
